#include "sweSolver.hpp"
#include "sweFunctions.hpp"
#include "exceptions.hpp"
#include "magic_enum.hpp"


namespace hms
{

/** @brief Default constructor, sets constants such as density and viscosity
 * to their default values. Also sets the default flux scheme to HLLC.
 * All other constructors must delegate to this one to make sure constants, etc.
 * are set correctly.
 */
template<typename MeshType>
SWESolver<MeshType>::SWESolver(){
	setConstants();
	/* do that here because it's a SolverBase member variable and SolverBase
	 * shouldn't get a default value that only applies to SWESolver */
	this->settings().fluxScheme = FluxScheme::HLLC;
}

/** @brief Constructs an SWESolver and resizes its fields according to the mesh
 * argument.
 */
template<typename MeshType>
SWESolver<MeshType>::SWESolver( const MeshBase<MeshType>& mesh ) :
	SWESolver()
{
	resizeFields(mesh);
}

/**
 * @brief Constructs an SWESolver by chaining to SWESolver(const MeshBase<MeshType>&)
 * and setting time and a constant uniform Manning's coefficient.
 * This ctor is meant for quick and easy construction, e.g. for unit testing.
 * With SolverBase::useSourceFields all off by default, only initialisation
 * of velocity, water depth, bottom elevation and their boundary conditions
 * is required to get the solver working.
 * 
 * @tparam MeshType 
 * @param mesh 
 * @param time 
 * @param uniformManningsN optional. Set to 0.025 by default.
 */
template<typename MeshType>
SWESolver<MeshType>::SWESolver(
	const MeshBase<MeshType>& mesh,
	const Time& time,
	scalar uniformManningsN
) : SWESolver(mesh)
{
	this->time( time );
	this->n() =
		ConstantField<Uniform, MeshType>{ mesh, ArrayNs<1>{uniformManningsN} };
}

/** @brief Sets the solver's constants to reasonable default values.
 */
template<typename MeshType>
void SWESolver<MeshType>::setConstants(){
	this->constant( ConstantName::Density ) = 1000; // kg/m3
	this->constant( ConstantName::Viscosity ) = 1e-6; // m2/s at 20 deg Celsius
}


namespace internal {

template<typename MeshType>
void resizeField(
	Field<MeshType>& field,
	const MeshBase<MeshType>& mesh,
	Index nComponents
){
	if ( !hasMatchingSize( field.values(), mesh, nComponents ) ){
		field = Field{ mesh, nComponents };
	}
}

} // end namespace internal

/** @brief Resizes the solver's fields according to the mesh argument.
 */
template<typename MeshType>
void SWESolver<MeshType>::resizeFields( const MeshBase<MeshType>& mesh ){
	using internal::resizeField;
	resizeField(v(), mesh, 2);
	resizeField(d(), mesh, 1);
	resizeField(z(), mesh, 1);
	resizeStorage();
	// resizeField(q(), mesh, 3);

	#ifdef SOLVER_STORE_FLUXES
	resizeField( this->fluxes(), mesh, 3);
	#endif

	#ifdef SWE_STORE_WATER_ELEVATION
	resizeField( h(), mesh, 1 );
	#endif
}

/** @brief Resizes the storage field in accordance with the mesh size. Does
 * nothing if the storage field already has the correct size */
template<typename MeshType>
void SWESolver<MeshType>::resizeStorage(){
	if ( !(this->settings().useIntegratedBlockLoop) ){
		this->logger().debug( "Resizing storage field..." );
		internal::resizeField( q(), this->mesh(), 3 );
	}
}

/** @brief Returns the density. */
template<typename MeshType>
auto SWESolver<MeshType>::rho() const -> scalar {
	return this->constant( ConstantName::Density );
}

/** @brief Sets the density to the provided value. */
template<typename MeshType>
void SWESolver<MeshType>:: rho( scalar s ){
	assert( s > 0 );
	this->constant( ConstantName::Density ) = s;
}

/** @brief Returns the viscosity. */
template<typename MeshType>
auto SWESolver<MeshType>::nu() const -> scalar {
	return this->constant( ConstantName::Viscosity );
}

/** @brief Sets the viscosity to the provided value. */
template<typename MeshType>
void SWESolver<MeshType>:: nu( scalar s){
	assert( s > 0 );
	this->constant( ConstantName::Viscosity ) = s;
}

/** @brief Returns a non-const reference to the velocity field. */
template<typename MeshType>
auto SWESolver<MeshType>::v() -> Field<MeshType>& {
	return this->field( FieldName::Velocity );
}

/** @brief Returns a const reference to the velocity field. */
template<typename MeshType>
auto SWESolver<MeshType>::v() const -> const Field<MeshType>& {
	return this->field( FieldName::Velocity );
}

/** @brief Returns a non-const reference to the water depth field. */
template<typename MeshType>
auto SWESolver<MeshType>::d() -> Field<MeshType>& {
	return this->field( FieldName::WaterDepth );
}

/** @brief Returns a const reference to the water depth field. */
template<typename MeshType>
auto SWESolver<MeshType>::d() const -> const Field<MeshType>& {
	return this->field( FieldName::WaterDepth );
}


#ifdef SWE_STORE_WATER_ELEVATION
/** @brief Returns a non-const reference to the water elevation field. */
template<typename MeshType>
auto SWESolver<MeshType>::h() -> Field<MeshType>& {
	return m_waterElevation;
}

/** @brief Returns a const reference to the water elevation field. */
template<typename MeshType>
auto SWESolver<MeshType>::h() const -> const Field<MeshType>& {
	return m_waterElevation;
}
#endif

/** @brief Returns a non-const reference to the storage field. */
template<typename MeshType>
auto SWESolver<MeshType>::q() -> Field<MeshType>& {
	return this->m_storage;
}

/** @brief Returns a const reference to the bottom elevation field. */
template<typename MeshType>
auto SWESolver<MeshType>::z() const -> const Field<MeshType>& {
	return this->field( FieldName::BottomElevation );
}

/** @brief Returns a non-const reference to the bottom elevation field. */
template<typename MeshType>
auto SWESolver<MeshType>::z() -> Field<MeshType>& {
	return this->field( FieldName::BottomElevation );
}

/** @brief Returns a non-const reference to the Manning's n field variant. */
template<typename MeshType>
auto SWESolver<MeshType>::n() -> ConstantFieldVar<MeshType>& {
	return this->constantField( ConstantFieldName::ManningsN );
}

/** @brief Returns a const reference to the Manning's n field variant. */
template<typename MeshType>
auto SWESolver<MeshType>::n() const -> const ConstantFieldVar<MeshType>& {
	return this->constantField( ConstantFieldName::ManningsN );
}

/** @brief Returns a non-const reference to the rainfall source field. Setting
 * this to positive values makes it act as a mass source. */
template<typename MeshType>
auto SWESolver<MeshType>::r() -> SourceFieldVar<MeshType>& {
	return this->sourceField( SourceFieldName::Rainfall );
}

/** @brief Returns a non-const reference to the infiltration source field.
 * Setting this to positive values makes it act as a mass sink. */
template<typename MeshType>
auto SWESolver<MeshType>::i() -> SourceFieldVar<MeshType>& {
	return this->sourceField( SourceFieldName::Infiltration );
}

/** @brief Returns true if infiltration source term calculation is enabled and
 * false otherwise */
template<typename MeshType>
bool SWESolver<MeshType>::useInfiltration() const {
	return this->useSourceField(SourceFieldName::Infiltration);
}

/** @brief Enable or disable infiltration source term calculation */
template<typename MeshType>
void SWESolver<MeshType>::useInfiltration( bool b ){
	this->useSourceField(SourceFieldName::Infiltration, b);
}

/** @brief Returns true if the rainfall source term is enabled and false
 * otherwise */
template<typename MeshType>
bool SWESolver<MeshType>::useRainfall() const {
	return this->useSourceField(SourceFieldName::Rainfall);
}

/** @brief Enable or disable rainfall source term usage */
template<typename MeshType>
void SWESolver<MeshType>::useRainfall( bool b ){
	this->useSourceField(SourceFieldName::Rainfall, b);
}



/** @brief Calculates the storage field from the values for water depth and
 * velocity. Only required once. */
template<typename MeshType>
void SWESolver<MeshType>::calculateStorage(){
	if ( !(this->settings().useIntegratedBlockLoop) ){
		this->logger().debug("Calculating storage field...");
		q().domainValues().row(0) = d().domainValues();
		q().domainValues().template bottomRows<2>() =
			v().domainValues() *
			d().domainValues().template replicate<2,1>();
	}
}

/** @brief If the mesh type is derived from StructMeshBase, sets the outside
 * corners of the field, which do not correspond to cells, to a value of zero. */
template<typename MeshType>
void SWESolver<MeshType>:: initialiseCorners(){
	this->logger().debug("Setting field corners to signal values...");
	setCorners( d().values(), this->mesh(), 0 );
}

/** @brief Calls all initialisation functions, i.e. those that are only supposed
 * to be executed once */
template<typename MeshType>
void SWESolver<MeshType>::initialiseImp(){
	this->logger().debug("Verifying matching boundary conditions...");
	for ( const BoundaryPatch& patch : this->mesh().boundaryPatches() )
		verifyMatchingBoundaryConditions(patch);

	/* ensure non-negative water depths,
	 * before setting corners and computing storage */
	d().values() = max( d().values(), 0 );

	initialiseCorners();
	resizeStorage();
	calculateStorage();

	#ifdef SOLVER_STORE_FLUXES
	internal::resizeField( this->fluxes(), this->mesh(), 3);
	#endif
}

namespace internal
{

template<Uniformity isUniform, typename MeshType>
bool manningsNOutOfRange( const ConstantField<isUniform, MeshType>& n ){
	if constexpr( isUniform ){
		return (n.storage() <= 0).any() || (n.storage() > 1).any();
	} else {
		return (n.domainValues() <= 0).any() || (n.domainValues() > 1).any();
	}
}

} // namespace internal


template<typename MeshType>
void SWESolver<MeshType>::validateImp() const {
	/* anything additional that needs to be checked before start, like
	 * value ranges etc. */
	std::visit( [&](const auto& n){
		if ( internal::manningsNOutOfRange(n) ){
			throw InputError( fmt::format(
				"ManningsN values out of range! Possible range: (0,1]. "
				"Values:\n{}", n.storage()
			) );
		}
	}, this->n() );
}


/** @brief Checks whether all field sizes match the mesh size */
template<typename MeshType>
void SWESolver<MeshType>:: checkFieldSizesImp() const {
	#ifndef NDEBUG

	assert( this->v().nComponents() == 2 );
	assert( this->d().nComponents() == 1 );
	if ( !(this->settings().useIntegratedBlockLoop) )
		assert( this->storage().nComponents() == 3 );
	auto test = [](const auto& field){
		assert( field.nComponents() == 1 );
	};
	for ( const auto& field : this->m_constantFields )
		std::visit(test, field);

	this->visitUsedSourceFields( test );

	#ifdef SOLVER_STORE_FLUXES
	assert( this->fluxes().nComponents() == 3 );
	#endif
	
	#endif
}



template class SWESolver<Mesh>;
template class SWESolver<UniMesh>;
template class SWESolver<RectMesh>;
template class SWESolver<StructMesh>;

} // end namespace hms